#!/usr/bin/env python
# -*- coding: UTF-8 -*-

from subprocess import Popen, PIPE
import json
import os
import time

def fetch_gpu_state():
    #  success  nvidia-smi命令是否可以查询到显卡状态 1为成功，0为失败
    success=0

    try:
        raw_data = Popen(['nvidia-smi'], stdout=PIPE, stderr=PIPE,shell=True).communicate()[0]
        string=raw_data.decode('UTF-8')
        if 'NVIDIA-SMI' in string and 'Driver Version' in string :
            success=1
        else:
            success=0

    except OSError:
        pass
    create_record(success)


def create_record(value):
    record = {}
    record['metric'] = "GPU.nvidia-smi"
    record['endpoint'] = os.uname()[1]
    record['timestamp'] = int(time.time())
    record['step'] = 60
    record['value'] = value
    record['counterType'] = 'GAUGE'
    record['tags'] = ''
    data.append(record)


if __name__ == '__main__':
    data = []
    fetch_gpu_state()
    print(json.dumps(data))
